#!/usr/bin/python

import sys
import bin_parser
import traceback
import re
import os
import struct
import yaml
import copy

from bin_parser import BinParser, BinParseFunctions

class ParseInfo(object):
    def __init__(self):
        self.clear()
        
    def loadData(self, data_file):
        self._inputHandle = open(data_file, 'rb')
        self._startAddr = 0
        self._dataLength = 0
        self._data = None
        self._inputHandle.seek(0, os.SEEK_END)
        self._length = self._inputHandle.tell()
        self._inputFile = data_file

    def loadTypes(self, types_file):
        tdata = yaml.load(open(types_file, 'rt'))
        # TODO join the new types
        self._types = tdata['types'] if 'types' in tdata else {}
        self._constants = tdata['constants'] if 'constants' in tdata else {}
        self._defaults = tdata['defaults'] if 'defaults' in tdata else {}

        # Add standard data types.
        self._types['raw'] = {}
        self._types['list'] = {}
        # Set default data type.
        if 'type' not in self._defaults:
            self._defaults['type'] = 'text'
        self._typesFiles.append(types_file)

    def loadStructure(self, struct_file):
        # TODO join the new structures
        self._structure = yaml.load(open(struct_file, 'rt'))
        self._structFiles.append(struct_file)

    def loadFunction(self, func_name, func_file):
        try:
            g = {}
            execfile(func_file, g)
            self._functions = g[func_name]
            print (self._functions)
        except:
            traceback.print_exc()
            return

    def info(self, out=sys.stdout):
        if not self._inputFile:
            return
        out.write('input file :%s, file length: %d\n' % (self._inputFile, self._length))
        out.write('types:\n')
        out.write('\t%r\n' % self._typesFiles)
        out.write('\t%r\n' % self._types)
        out.write('constants:\n')
        out.write('\t%r\n' % self._constants)
        out.write('defaults:\n')
        out.write('\t%r\n' % self._defaults)
        out.write('structure:\n')
        out.write('\t%r\n' % self._structFiles)
        out.write('\t%r\n' % self._structure)
        out.write('functions:\n')
        out.write('\t%r\n' % self._functions)
        
    def clear(self):
        self._startAddr = 0
        self._dataLength = 0
        self._data = None
        self._inputHandle = None
        self._types = None
        self._structure = None
        self._inputFile = None
        self._typesFiles = []
        self._structFiles = []
        self._functions = BinParseFunctions

    def getData(self, startAddr, length):
        if self._data != None and startAddr == self._startAddr and length == self._dataLength:
            return self._data

        if not self._inputHandle:
            print ("input handle is not opened")
            return None

        if self._length <= startAddr:
            print ("start address is too large:%s, file length is %s" % (startAddr, self._length))
            return None

        if length == 0 or self._length < (startAddr + length):
            length = self._length - startAddr

        self._inputHandle.seek(startAddr, os.SEEK_SET)
        self._data = self._inputHandle.read(length)
        self._startAddr = startAddr
        self._dataLength = length
        return self._data


    def storeValue(self, dest, ret_names, ret):
        scope = copy.copy(ret)
        for ret_name in ret_names:
            if type(ret_name) == list:
                name = ret_name[0]
                value = ret_name[1]
            else:
                name = ret_name
                value = None

            if not name:
                name = value

            if not value:
                ret_value = ret 
            else:
                try:
                    ret_value = eval(value, scope)
                except:
                    traceback.print_exc()
                if not ret_value:
                    print ("try get the value failed:%s" % value)
                    continue
                
            if name in dest:
                target = dest[name]
                if type(target) != list:
                    target = [target, ret_value]
                else:
                    target.append(ret_value)
                dest[name] = target
            else:
                dest[name] = ret_value


    def parse(self, startAddr, length, struct_name, loop = False, ret_names = None, dests=None, out=sys.stdout):
        data = self.getData(startAddr, length)
        length = len(data)
        if not data:
            print ("cannot get the data")
            return
        structure = self.findStructure(struct_name)
        if not structure:
            print ("cannot find structure : %s"% struct_name )
            return
        #print(struct_name, structure)
        parser = bin_parser.BinParser()
        count = 0
        dest = []
        while count < length:
            size = parser.parse(data[count:], self._types, self._constants, self._defaults, structure, self._functions if self._functions else BinParseFunctions)
            if size <= 0:
                break
            dest.append(parser.parsed)
            if out:
                parser.write(out)
            count += size
            if not loop:
                break

        if dests == None:
            return

        if not ret_names:
            ret_names = '_last_result'
        
        if type(ret_names) == list:
            for ret in dest:
                self.storeValue(dests, ret_names, ret)
        else:
            if loop:
                dests[ret_names] = dest
            else:
                dests[ret_names] = dest[0]


    def findStructure(self, struct_name):
        if not struct_name or struct_name == '*':
            return self._structure
        queue = [self._structure]
        while len(queue) > 0:
            structure = queue[0]
            queue = queue[1:]
            for item in structure:
                if 'name' in item and item['name'] == struct_name:
                    return [item]
                elif 'structure' in item:
                    queue.append(item['structure'])
        return None
            

#show command

PARSE_INFO = ParseInfo()
USER_VAR = {}
USER_COMMAND = None

print ("parse bin file:")

def do_help(args):
    print('''
    help : show help;
    exit : exit this command;
    open <file name> [<struct> <types>]: open the special file;
    struct <struct.yml>: load the structer file;
    types  <types.yml> : load the types file;
    dump <format> <start addr> <length>: dump the special range,
        format: <word-width><type>, word-width is '1','2','4','8', by byte,
            type is 'x': hex, 'o': oct, 'b': binary.
        eg: dump 4x 0x20000 0x100
    parse <struct>:<start addr>[#<length>] [loop] [=<var_name>]
    p  <var_name or var_exception>
    info : show the parse info
    script: load and execute a script file
    function <funcname> <file or code>: load a file and execute, get the funcname
        eg: function Foo foo.py   load foo.py as python code file, and get the class Foo as function
    usercommand <command_set_name> <command_define_file>: load a user defined command
        user command use '.' + command
    ''')


def do_open(args):
    global PARSE_INFO
    if len(args) <= 1:
        help([])
        return
    PARSE_INFO.loadData(args[1])
    if len(args) >= 3:
        PARSE_INFO.loadStructure(args[2])
    if len(args) >= 4:
        PARSE_INFO.loadTypes(args[3])

def do_struct(args):
    global PARSE_INFO
    if len(args) <= 1:
        help([])
        return
    PARSE_INFO.loadStructure(args[1])
    

def do_types(args):
    global PARSE_INFO
    if len(args) <= 1:
        help([])
        return
    PARSE_INFO.loadTypes(args[1])


def dump_data(start, length, word_width, byte_format):
    global PARSE_INFO
    data = PARSE_INFO.getData(start, length) 
    length = len(data)
    if byte_format == 'x':
        str_format = '%%0%dx' % (word_width * 2)
        word_index = [1,2,4,8].index(word_width)
        pack_format = 'BHIQ'[word_index:word_index+1]
        #print (str_format, pack_format, word_width)

        for i in range(0, length, 8):
            sys.stdout.write ("%08X: "% (start + i))
            end = i + 8 
            if end > length: end = length
            last = end % word_width
            end -= last
            s = ""
            for j in range (i, end, word_width):
                v = struct.unpack(pack_format, data[j:j + word_width])
                s += ' ' + str_format % v
            for j in range (end, end + last):
                v = struct.unpack('B', data[j:j+1])
                s += ' %02X' % v
        
            print(s)
            

    elif byte_format == 'o':
        pass
    elif byte_format == 'b':
        pass



def do_dump(args):
    if len(args) < 3:
        help([])
        return
    args = args[1:]
    if len(args) >= 3:
        format = args[0]
        args = args[1:]
    else:
        format = "2x"

    try:
        start = eval(args[0])
        length = eval(args[1])
    except:
        traceback.print_exc()
        return
        

    word_width = int(format[0])
    byte_format = format[1]
    if not word_width in [1,2,4,8]:
        word_width = 2
    if not byte_format in ['b','x', 'o']:
        byte_format = 'x'

    dump_data(start, length, word_width, byte_format) 

def do_parse(args):
    global PARSE_INFO
    global USER_VAR
    if len(args) < 2:
        help([])
        return

    start = 0
    length = 0
    loop = False
    ret_names = None
    struct_name = None
    expr_start = None
    expr_length = None

    for arg in args[1:]:
        if arg == 'loop':
            loop = True
        elif arg.find('=') >= 0:
            ret_names = []
            for expr in arg.split(','):
                ret_names.append(expr.split('='))
                
        else:
            idx = arg.find(':')
            if idx > 0:
                struct_name = arg[0:idx]
                arg = arg[idx+1:]
                idx = arg.find('#')
                if idx < 0:
                    expr_start = arg
                else:
                    expr_start = arg[0:idx] 
                    expr_length = arg[idx+1:]
            else:
                print ("Unkown parse command: %s" % arg)

    try:
        scope = copy.copy(USER_VAR)
        if expr_start:
            start = eval(expr_start, scope)
        if expr_length:
            length = eval(expr_length, scope)
    except:
        traceback.print_exc()
        return
 
    PARSE_INFO.parse(start, length, struct_name, loop, ret_names, USER_VAR);

def do_info(args):
    global PARSE_INFO
    PARSE_INFO.info()

def do_clear(args):
    global PARSE_INFO
    PARSE_INFO.clear()

def do_echo(args):
    print (' '.join(args[1:]))

def do_script(args):
    if len(args) <= 1:
        print ("Please give the script file")
    do_command_from_file(args[1])

def do_function(args):
    global PARSE_INFO
    if len(args) <= 2:
        print ("Please give the function file")
    PARSE_INFO.loadFunction(args[1], args[2]) 

def do_usercommand(args):
    global USER_COMMAND
    global USER_VAR
    global PARSE_INFO
    if len(args) <= 2:
        print ("Please give the user defined file")
        help([])
        return

    try:
        g = {}
        execfile(args[2], g)
        func = g[args[1]]
        USER_COMMAND = func(PARSE_INFO, USER_VAR)
    except:
        traceback.print_exc()

# do print USER_VAR
def do_p(args):
    global USER_VAR
    if len(args) <= 1:
        print (USER_VAR)
        return

    if args[1] == '?':
        expr = '_last_result'
    else:
        expr = args[1]
    try:
        scope = copy.copy(USER_VAR)
        ret = eval(expr, scope)
        print (ret)
    except:
        traceback.print_exc()

################################################################
##
def do_command(cmd, in_stdin):
    global USER_COMMAND
    global USER_VAR
    if len(cmd) == 0:
        return 0
    cmd_list = re.split(' +', cmd)
    #print (cmd_list)
    if cmd_list[0] == 'exit' or cmd_list[0] == 'q':
        return -1
    if cmd_list[0] == 'stdin' and not in_stdin:
        do_command_from_stdin()
        return 0

    if cmd_list[0][0] == '.': # user command
        if not USER_COMMAND:
            print ("User Command is not defined!")
            return -2
        cmd_list[0] = cmd_list[0][1:] 
        try:
            # try evaluate paramters first
            scope = copy.copy(USER_VAR)
            try:
                user_scope = getattr(USER_COMMAND, 'getLocalScope')()
                scope.extend(user_scope)
            except:
                pass

            args = []
            for cmd in cmd_list[1:]:
                if ord(cmd[0]) in range(49,58) or cmd[0] == '$':
                    if len(cmd) >= 2 and cmd[1] == '(' and cmd[-1] == ')':
                        expr = cmd[2:-1]
                    else:
                        expr = cmd
                    try:
                        # eval the args
                        value = eval(expr, scope)
                        args.append(value)
                    except:
                        traceback.print_exc()
                        args.append(cmd)
                else:
                    args.append(cmd)
                        
            ret = getattr(USER_COMMAND, cmd_list[0])(*args)
            print (ret) 
        except:
            traceback.print_exc()
            print("command execute %s in %r faield" % (cmd_list[0], USER_COMMAND))
            return -2
        
        return 0

    try:
        func_name = eval('do_' + cmd_list[0])
        func_name(cmd_list)
    except:
        print("unkown command name :" + cmd_list[0])
        traceback.print_exc()
        help(cmd_list)
        return -2
    return 0
    


def do_command_from_stdin():
    while True:
        cmd_str = raw_input("> ")
        cmd_str = cmd_str.strip()
        if do_command(cmd_str, True) == -1: #exit
            break
    
def do_command_from_file(file_name):
    f = open(file_name, 'r')
    if not f:
        print ("Cannot open script file : %s" % file_name)
        return
    while True:
        line = f.readline()
        if line:
            line = line.strip()
            if len(line) == 0 or line[0] == '#':
                continue
            if do_command(line, False) != 0:
                break
        else:
            break
    f.close()



#########################################################
def main(argv):
    if len(argv) > 1:
        do_command_from_file(argv[1])
    else:
        do_command_from_stdin()


main(sys.argv)

